# Transformations module for FLEXINVERT.jl
# Ports mod_transform.F/.f90 functionality for phi↔chi transforms

module Transformations

using LinearAlgebra
using Printf: @sprintf
using ..PriorCovariance: CovarianceMatrix

export phi2chi, chi2phi, phi2chi!, chi2phi!, apply_transform, apply_adjoint_transform,
       apply_B_inv, build_H_reduced_with_basis, build_H_reduced, build_zwork,
       apply_hbht_column_scaling,
       full_covariance_matrix, full_precision_matrix

"""
    phi2chi(phi::Vector{Float64}, cov::CovarianceMatrix) -> Vector{Float64}

Transform control variables (phi) to physical variables (chi).
phi is in the reduced eigenspace, chi is in physical space.

The transformation is: chi = U * sqrt(Λ) * phi
where U are eigenvectors and Λ are eigenvalues of the covariance matrix.
"""
function phi2chi(phi::Vector{Float64}, cov::CovarianceMatrix)
    if length(phi) != cov.n_modes
        error("phi vector length ($(length(phi))) must match number of modes ($(cov.n_modes))")
    end

    # chi = U * sqrt(Λ) * phi
    sqrt_eigenvals = sqrt.(cov.eigenvalues)
    weighted_phi = sqrt_eigenvals .* phi
    chi = cov.eigenvectors * weighted_phi
    if cov.has_numscale_area_scaling
        chi = cov.numscale_area_scaling .* chi
    end
    if cov.has_physical_scaling
        chi = cov.physical_scaling .* chi
    end
    if cov.has_row_scaling
        chi = cov.row_scaling .* chi
    end

    return chi
end

"""
    chi2phi(chi::Vector{Float64}, cov::CovarianceMatrix) -> Vector{Float64}

Transform physical variables (chi) to control variables (phi).
chi is in physical space, phi is in the reduced eigenspace.

The transformation is: phi = (1/sqrt(Λ)) * U^T * chi
where U are eigenvectors and Λ are eigenvalues of the covariance matrix.
"""
function chi2phi(chi::Vector{Float64}, cov::CovarianceMatrix)
    expected_size = cov.nbox * cov.nt
    if length(chi) != expected_size
        error("chi vector length ($(length(chi))) must match physical space size ($expected_size)")
    end

    scaled_chi = chi
    if cov.has_row_scaling
        scaled_chi = scaled_chi .* cov.inv_row_scaling
    end
    if cov.has_physical_scaling
        scaled_chi .*= cov.inv_physical_scaling
    end
    if cov.has_numscale_area_scaling
        scaled_chi .*= cov.inv_numscale_area_scaling
    end

    # phi = (1/sqrt(Λ)) * U^T * chi
    proj_chi = cov.eigenvectors' * scaled_chi
    sqrt_eigenvals = sqrt.(cov.eigenvalues)
    phi = proj_chi ./ sqrt_eigenvals

    return phi
end

"""
    phi2chi!(chi::Vector{Float64}, phi::Vector{Float64}, cov::CovarianceMatrix)

In-place transformation from control variables (phi) to physical variables (chi).
"""
function phi2chi!(chi::Vector{Float64}, phi::Vector{Float64}, cov::CovarianceMatrix)
    if length(phi) != cov.n_modes
        error("phi vector length ($(length(phi))) must match number of modes ($(cov.n_modes))")
    end

    expected_size = cov.nbox * cov.nt
    if length(chi) != expected_size
        error("chi vector length ($(length(chi))) must match physical space size ($expected_size)")
    end

    # chi = U * sqrt(Λ) * phi
    sqrt_eigenvals = sqrt.(cov.eigenvalues)
    weighted_phi = sqrt_eigenvals .* phi
    mul!(chi, cov.eigenvectors, weighted_phi)
    if cov.has_numscale_area_scaling
        chi .*= cov.numscale_area_scaling
    end
    if cov.has_physical_scaling
        chi .*= cov.physical_scaling
    end
    if cov.has_row_scaling
        chi .*= cov.row_scaling
    end

    return chi
end

"""
    chi2phi!(phi::Vector{Float64}, chi::Vector{Float64}, cov::CovarianceMatrix)

In-place transformation from physical variables (chi) to control variables (phi).
"""
function chi2phi!(phi::Vector{Float64}, chi::Vector{Float64}, cov::CovarianceMatrix)
    expected_size = cov.nbox * cov.nt
    if length(chi) != expected_size
        error("chi vector length ($(length(chi))) must match physical space size ($expected_size)")
    end

    if length(phi) != cov.n_modes
        error("phi vector length ($(length(phi))) must match number of modes ($(cov.n_modes))")
    end

    scaled_chi = chi
    if cov.has_row_scaling
        scaled_chi = scaled_chi .* cov.inv_row_scaling
    end
    if cov.has_physical_scaling
        scaled_chi .*= cov.inv_physical_scaling
    end
    if cov.has_numscale_area_scaling
        scaled_chi .*= cov.inv_numscale_area_scaling
    end

    # phi = (1/sqrt(Λ)) * U^T * chi
    # Use temporary vector to avoid allocation
    temp = cov.eigenvectors' * scaled_chi
    sqrt_eigenvals = sqrt.(cov.eigenvalues)
    phi .= temp ./ sqrt_eigenvals

    return phi
end

"""
    apply_transform(x::Vector{Float64}, cov::CovarianceMatrix, forward::Bool=true) -> Vector{Float64}

Apply transformation in forward (phi→chi) or adjoint (chi→phi) direction.

# Arguments
- `x`: Input vector
- `cov`: Covariance matrix with eigenvalue decomposition
- `forward`: If true, apply phi→chi; if false, apply chi→phi

# Returns
- Transformed vector
"""
function apply_transform(x::Vector{Float64}, cov::CovarianceMatrix, forward::Bool=true)
    if forward
        return phi2chi(x, cov)
    else
        return chi2phi(x, cov)
    end
end

"""
    apply_adjoint_transform(x::Vector{Float64}, cov::CovarianceMatrix, forward::Bool=true) -> Vector{Float64}

Apply adjoint transformation.

# Arguments
- `x`: Input vector
- `cov`: Covariance matrix with eigenvalue decomposition
- `forward`: If true, apply adjoint of phi→chi (i.e., chi→phi); if false, apply adjoint of chi→phi

# Returns
- Adjoint-transformed vector
"""
function apply_adjoint_transform(x::Vector{Float64}, cov::CovarianceMatrix, forward::Bool=true)
    # Adjoint of phi→chi is chi→phi and vice versa
    return apply_transform(x, cov, !forward)
end

"""
    transform_gradient(grad_chi::Vector{Float64}, cov::CovarianceMatrix) -> Vector{Float64}

Transform gradient from physical space (chi) to control space (phi).
Used in optimization where gradient is computed in physical space but optimization occurs in control space.

The transformation is: grad_phi = (1/sqrt(Λ)) * U^T * grad_chi
"""
function transform_gradient(grad_chi::Vector{Float64}, cov::CovarianceMatrix)
    return chi2phi(grad_chi, cov)
end

"""
    transform_state_increment(dphi::Vector{Float64}, cov::CovarianceMatrix) -> Vector{Float64}

Transform state increment from control space (phi) to physical space (chi).
Used in optimization to convert control space increments to physical space updates.

The transformation is: dchi = U * sqrt(Λ) * dphi
"""
function transform_state_increment(dphi::Vector{Float64}, cov::CovarianceMatrix)
    return phi2chi(dphi, cov)
end

"""
    build_H_reduced_with_basis(H, cov) -> (H_reduced, B_half)

Construct the reduced observation operator H · B^{1/2} and return the basis
matrix used (physical columns of B^{1/2}).
"""
function build_H_reduced_with_basis(H::AbstractMatrix, cov::CovarianceMatrix)
    n_modes = cov.n_modes

    # Cholesky-based basis (used internally for HBHᵀ solves)
    B_full = full_covariance_matrix(cov)
    chol = cholesky(Symmetric(B_full); check=false)
    basis_full = Matrix(chol.L)
    ncols = min(size(basis_full, 2), n_modes)
    basis = basis_full[:, 1:ncols]

    # Fortran-style basis (apply phi->chi transform to canonical basis vectors)
    n_physical = cov.nbox * cov.nt
    basis_fortran = Array{Float64}(undef, n_physical, n_modes)
    tmp_phi = zeros(Float64, n_modes)
    tmp_chi = zeros(Float64, n_physical)
    for i in 1:n_modes
        tmp_phi .= 0.0
        tmp_phi[i] = 1.0
        phi2chi!(tmp_chi, tmp_phi, cov)
        basis_fortran[:, i] = tmp_chi
    end

    map_chol_to_eig = basis_fortran \ basis

    output_dir = get(ENV, "FLEXINVERT_OUTPUT_DIR", "")
    if !isempty(output_dir)
        try
            open(joinpath(output_dir, "basis_cholesky.txt"), "w") do io
                for i in 1:size(basis, 1)
                    for j in 1:size(basis, 2)
                        j > 1 && write(io, ' ')
                        print(io, @sprintf("%.8e", basis[i, j]))
                    end
                    write(io, '\n')
                end
            end
            open(joinpath(output_dir, "basis_fortran.txt"), "w") do io
                for i in 1:size(basis_fortran, 1)
                    for j in 1:size(basis_fortran, 2)
                        j > 1 && write(io, ' ')
                        print(io, @sprintf("%.8e", basis_fortran[i, j]))
                    end
                    write(io, '\n')
                end
            end
        catch err
            @warn "Failed to write basis diagnostics" output_dir err
        end
    end

    H_reduced = H * basis
    return H_reduced, basis, map_chol_to_eig
end

"""
    build_H_reduced(H, cov) -> Matrix

Convenience wrapper returning only the reduced observation operator.
"""
function build_H_reduced(H::AbstractMatrix, cov::CovarianceMatrix)
    H_reduced, _, _ = build_H_reduced_with_basis(H, cov)
    return H_reduced
end

"""
    apply_hbht_column_scaling(H, cov) -> Matrix

Scale the observation operator columns using the HBHᵀ row-scaling factors
imported from the Fortran artefacts. Returns `H` unchanged when no scaling
is present.
"""
function apply_hbht_column_scaling(H::AbstractMatrix, cov::CovarianceMatrix)
    if !cov.has_hbht_row_scaling
        return H
    end
    ncols = size(H, 2)
    expected = length(cov.hbht_row_scaling)
    ncols == expected || error("HBHᵀ scaling length mismatch: matrix columns $ncols vs vector $expected")
    return H .* reshape(cov.hbht_row_scaling, 1, :)
end

function compute_fortran_covariance_block(cov::CovarianceMatrix)
    corr = cov.fortran_correlation
    xerr = cov.fortran_xerr
    corr === nothing && return nothing
    xerr === nothing && return nothing
    nbox = size(corr, 1)
    ndt = max(cov.ndt, 1)
    ndvar = nbox * ndt
    length(xerr) == ndvar || return nothing

    area = cov.fortran_area
    covmat = zeros(Float64, ndvar, ndvar)
    covsum = 0.0
    for nt in 0:ndt-1
        offset = nt * nbox
        for j in 1:nbox
            idx_j = offset + j
            xerr_j = xerr[idx_j]
            area_j = area !== nothing && j <= length(area) ? area[j] : 0.0
            for i in 1:nbox
                idx_i = offset + i
                val = corr[i, j] * xerr[idx_i] * xerr_j
                covmat[idx_i, idx_j] = val
                area_i = area !== nothing && i <= length(area) ? area[i] : 0.0
                covsum += val * area_i * area_j
            end
        end
    end

    covsum /= (ndt > 0 ? ndt^2 : 1.0)
    if cov.globerr > 0 && covsum > 0
        toterr = sqrt(covsum) * 3600.0 * 24.0 * 365.0 / 1e9 / NUMERICAL_SCALE
        if toterr > 0
            scale = cov.globerr / toterr
            covmat .*= scale^2
        end
    end
    covmat
end

"""
    build_zwork(H, cov) -> Matrix

Construct BHᵀ using the current covariance representation.
"""
function build_zwork(H::AbstractMatrix, cov::CovarianceMatrix)
    B_full = full_covariance_matrix(cov)
    z = B_full * transpose(H)
    return z
end

"""
    apply_B_inv(x::Vector{Float64}, cov::CovarianceMatrix) -> Vector{Float64}

Apply the inverse prior covariance matrix B⁻¹ to a physical-space vector.
"""
function apply_B_inv(x::Vector{Float64}, cov::CovarianceMatrix)
    expected = cov.nbox * cov.nt
    length(x) == expected || error("Input vector length $(length(x)) must match physical size ($expected)")

    temp = x
    if cov.has_row_scaling
        temp = temp .* cov.inv_row_scaling
    end
    if cov.has_physical_scaling
        temp .*= cov.inv_physical_scaling
    end
    if cov.has_numscale_area_scaling
        temp .*= cov.inv_numscale_area_scaling
    end
    temp = cov.eigenvectors' * temp
    temp .= temp ./ cov.eigenvalues
    result = cov.eigenvectors * temp
    if cov.has_numscale_area_scaling
        result .*= cov.inv_numscale_area_scaling
    end
    if cov.has_physical_scaling
        result .*= cov.inv_physical_scaling
    end
    if cov.has_row_scaling
        result .*= cov.inv_row_scaling
    end
    return result
end

"""
    chi2phi_fortran(phi::Vector{Float64}, cov::CovarianceMatrix, state_prior::Vector{Float64}) -> Vector{Float64}

Transform from control space (phi) to physical space (chi) using Fortran's algorithm.
This matches the Fortran mod_transform.f90:phi2chi function with temporal correlation matrices.

Arguments:
- phi: Control space vector
- cov: Covariance matrix with temporal correlation matrices
- state_prior: Prior state vector (x₀)

Returns:
- chi: Physical space vector (B^(-1/2) * (x - x₀))
"""
function chi2phi_fortran(phi::Vector{Float64}, cov::CovarianceMatrix, state_prior::Vector{Float64})
    nbox = cov.nbox
    ntstate = Int(cov.nt / max(1, Int(size(cov.temporal_corr, 1) > 1 ? 1 : cov.nt)))
    ndt = Int(cov.nt / ntstate)

    # Convert to multi-dimensional arrays for temporal processing
    phi_reshaped = reshape(phi, nbox, ntstate, ndt)
    chi_reshaped = zeros(Float64, nbox, ntstate, ndt)

    # Apply Fortran-style transformation with temporal correlations
    for jt in 1:ntstate
        for k in 1:nbox
            # Build intermediate vector (matches Fortran's tmp calculation)
            tmp = zeros(Float64, cov.n_modes)
            for j in 1:cov.n_modes
                tmp[j] = cov.eigenvectors[k, j] / sqrt(cov.eigenvalues[j])
            end

            # Apply temporal correlation (matches Fortran's isqcort)
            for it in 1:ntstate
                for dt in 1:ndt
                    source_idx = (it - 1) * ndt + dt
                    target_idx = (jt - 1) * ndt + dt

                    if source_idx <= size(phi_reshaped, 3) && target_idx <= size(chi_reshaped, 3)
                        # Apply temporal correlation matrix
                        for it_inner in 1:ntstate
                            corr_weight = cov.temporal_corr_sqrt_inv[jt, it_inner]
                            if abs(corr_weight) > 1e-15
                                inner_sum = 0.0
                                for j in 1:cov.n_modes
                                    inner_sum += tmp[j] * cov.eigenvectors[k, j]
                                end
                                chi_reshaped[k, jt, target_idx] += corr_weight * inner_sum *
                                    (phi_reshaped[k, it_inner, source_idx] -
                                     state_prior[(it_inner-1)*nbox + k])
                            end
                        end
                    end
                end
            end
        end
    end

    return reshape(chi_reshaped, nbox * ntstate * ndt)
end

"""
    phi2chi_fortran(chi::Vector{Float64}, cov::CovarianceMatrix, state_prior::Vector{Float64}) -> Vector{Float64}

Transform from physical space (chi) to control space (phi) using Fortran's algorithm.
This matches the Fortran mod_transform.f90:chi2phi function with temporal correlation matrices.

Arguments:
- chi: Physical space vector
- cov: Covariance matrix with temporal correlation matrices
- state_prior: Prior state vector (x₀)

Returns:
- phi: Control space vector
"""
function phi2chi_fortran(chi::Vector{Float64}, cov::CovarianceMatrix, state_prior::Vector{Float64})
    nbox = cov.nbox
    ntstate = Int(cov.nt / max(1, Int(size(cov.temporal_corr, 1) > 1 ? 1 : cov.nt)))
    ndt = Int(cov.nt / ntstate)

    # Convert to multi-dimensional arrays for temporal processing
    chi_reshaped = reshape(chi, nbox, ntstate, ndt)
    phi_reshaped = zeros(Float64, nbox, ntstate, ndt)

    # Apply Fortran-style transformation with temporal correlations
    for jt in 1:ntstate
        for k in 1:nbox
            # Build intermediate vector (matches Fortran's tmp calculation)
            tmp = zeros(Float64, cov.n_modes)
            for j in 1:cov.n_modes
                tmp[j] = cov.eigenvectors[k, j] * sqrt(cov.eigenvalues[j])
            end

            # Apply temporal correlation (matches Fortran's sqcort)
            for it in 1:ntstate
                for dt in 1:ndt
                    source_idx = (it - 1) * ndt + dt
                    target_idx = (jt - 1) * ndt + dt

                    if source_idx <= size(chi_reshaped, 3) && target_idx <= size(phi_reshaped, 3)
                        # Apply temporal correlation matrix
                        for it_inner in 1:ntstate
                            corr_weight = cov.temporal_corr_sqrt[jt, it_inner]
                            if abs(corr_weight) > 1e-15
                                inner_sum = 0.0
                                for j in 1:cov.n_modes
                                    inner_sum += tmp[j] * cov.eigenvectors[k, j]
                                end
                                phi_reshaped[k, jt, target_idx] += corr_weight * inner_sum *
                                    chi_reshaped[k, it_inner, source_idx]
                            end
                        end
                    end
                end
            end
        end
    end

    # Add prior state contribution
    result = reshape(phi_reshaped, nbox * ntstate * ndt) + state_prior

    return result
end

"""
    prod_icov_fortran(x::Vector{Float64}, cov::CovarianceMatrix) -> Vector{Float64}

Apply the inverse covariance matrix B^(-1) to vector x using Fortran's algorithm.
This matches the Fortran mod_transform.f90:prod_icov function.

Arguments:
- x: Input vector in physical space
- cov: Covariance matrix with temporal correlation matrices

Returns:
- result: B^(-1) * x
"""
function prod_icov_fortran(x::Vector{Float64}, cov::CovarianceMatrix)
    nbox = cov.nbox
    ntstate = Int(cov.nt / max(1, Int(size(cov.temporal_corr, 1) > 1 ? 1 : cov.nt)))
    ndt = Int(cov.nt / ntstate)

    # Convert to multi-dimensional arrays for temporal processing
    x_reshaped = reshape(x, nbox, ntstate, ndt)
    result_reshaped = zeros(Float64, nbox, ntstate, ndt)

    # Apply Fortran-style B^(-1) * x with temporal correlations
    for jt in 1:ntstate
        for k in 1:nbox
            # Build intermediate vector (matches Fortran's tmp calculation for B^(-1))
            tmp = zeros(Float64, cov.n_modes)
            for j in 1:cov.n_modes
                tmp[j] = cov.eigenvectors[k, j] / cov.eigenvalues[j]
            end

            # Apply temporal correlation (matches Fortran's icort)
            for it in 1:ntstate
                for dt in 1:ndt
                    source_idx = (it - 1) * ndt + dt
                    target_idx = (jt - 1) * ndt + dt

                    if source_idx <= size(x_reshaped, 3) && target_idx <= size(result_reshaped, 3)
                        # Apply temporal correlation matrix
                        for it_inner in 1:ntstate
                            corr_weight = cov.temporal_corr_inv[jt, it_inner]
                            if abs(corr_weight) > 1e-15
                                inner_sum = 0.0
                                for j in 1:cov.n_modes
                                    inner_sum += tmp[j] * cov.eigenvectors[k, j]
                                end
                                result_reshaped[k, jt, target_idx] += corr_weight * inner_sum *
                                    x_reshaped[k, it_inner, source_idx]
                            end
                        end
                    end
                end
            end
        end
    end

    return reshape(result_reshaped, nbox * ntstate * ndt)
end

"""
    validate_transforms(cov::CovarianceMatrix; rtol::Float64=1e-10) -> Bool

Validate that the forward and adjoint transforms are consistent.

# Arguments
- `cov`: Covariance matrix with eigenvalue decomposition
- `rtol`: Relative tolerance for validation

# Returns
- `true` if transforms are consistent, `false` otherwise
"""
function validate_transforms(cov::CovarianceMatrix; rtol::Float64=1e-10)
    @info "Validating transform consistency"

    # Test round-trip in control space
    phi_test = randn(cov.n_modes)
    chi_intermediate = phi2chi(phi_test, cov)
    phi_recovered = chi2phi(chi_intermediate, cov)

    phi_error = norm(phi_test - phi_recovered) / norm(phi_test)
    @info "Control space round-trip relative error: $phi_error"

    if phi_error > rtol
        @warn "Control space round-trip error ($phi_error) exceeds tolerance ($rtol)"
        return false
    end

    # Test orthogonality of eigenvectors
    U = cov.eigenvectors
    orthogonality_error = norm(U' * U - I)
    @info "Eigenvector orthogonality error: $orthogonality_error"

    if orthogonality_error > rtol
        @warn "Eigenvector orthogonality error ($orthogonality_error) exceeds tolerance ($rtol)"
        return false
    end

    # Test that eigenvalues are positive
    min_eigenval = minimum(cov.eigenvalues)
    @info "Minimum eigenvalue: $min_eigenval"

    if min_eigenval <= 0
        @warn "Non-positive eigenvalue detected: $min_eigenval"
        return false
    end

    @info "Transform validation passed"
    return true
end

"""
    effective_rank(cov::CovarianceMatrix; threshold::Float64=1e-12) -> Int

Calculate the effective rank of the covariance matrix.

# Arguments
- `cov`: Covariance matrix with eigenvalue decomposition
- `threshold`: Eigenvalue threshold for determining rank

# Returns
- Effective rank (number of eigenvalues above threshold)
"""
function effective_rank(cov::CovarianceMatrix; threshold::Float64=1e-12)
    return count(λ -> λ > threshold, cov.eigenvalues)
end

"""
    condition_number(cov::CovarianceMatrix) -> Float64

Calculate the condition number of the covariance matrix.

# Returns
- Condition number (ratio of largest to smallest retained eigenvalue)
"""
function condition_number(cov::CovarianceMatrix)
    λ_max = maximum(cov.eigenvalues)
    λ_min = minimum(cov.eigenvalues[cov.eigenvalues .> 1e-12])
    return λ_max / λ_min
end

function full_covariance_matrix(cov::CovarianceMatrix)
    if cov.full_matrix !== nothing
        return Matrix(cov.full_matrix)
    end
    B = cov.eigenvectors * Diagonal(cov.eigenvalues) * cov.eigenvectors'
    if cov.has_numscale_area_scaling
        N = Diagonal(cov.numscale_area_scaling)
        B = N * B * N
    end
    if cov.has_physical_scaling
        D = Diagonal(cov.physical_scaling)
        B = D * B * D
    end
    if cov.has_row_scaling
        R = Diagonal(cov.row_scaling)
        B = R * B * R
    end
    return B
end

function full_precision_matrix(cov::CovarianceMatrix)
    if cov.full_matrix !== nothing
        return inv(Matrix(cov.full_matrix))
    end
    B_inv = cov.eigenvectors * Diagonal(1.0 ./ cov.eigenvalues) * cov.eigenvectors'
    if cov.has_numscale_area_scaling
        N_inv = Diagonal(cov.inv_numscale_area_scaling)
        B_inv = N_inv * B_inv * N_inv
    end
    if cov.has_physical_scaling
        D_inv = Diagonal(cov.inv_physical_scaling)
        B_inv = D_inv * B_inv * D_inv
    end
    if cov.has_row_scaling
        R_inv = Diagonal(cov.inv_row_scaling)
        B_inv = R_inv * B_inv * R_inv
    end
    return B_inv
end

end # module Transformations
